# -*- coding: utf-8 -*-
"""
请求体缓存中间件

解决FastAPI中请求体只能读取一次的问题：
- 在ASGI层面提前读取并缓存请求体
- 创建可重复读取的receive函数
- 允许多个中间件和路由处理器访问请求体
- 支持分块传输和大文件限制
"""

import asyncio
import logging
from typing import Callable, Dict, Any, Awaitable, Tuple
from starlette.types import Scope, Receive, Send, Message

logger = logging.getLogger(__name__)


class RequestBodyCacheMiddleware:
    """请求体缓存中间件 - 在ASGI层面解决请求体只能读取一次的问题"""
    
    def __init__(
        self, 
        app: Callable[[Scope, Receive, Send], Awaitable[None]],
        max_body_size: int = 10 * 1024 * 1024  # 10MB
    ):
        """
        初始化请求体缓存中间件
        
        Args:
            app: ASGI应用
            max_body_size: 最大请求体大小（字节）
        """
        self.app = app
        self.max_body_size = max_body_size
    
    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """ASGI调用入口"""
        
        # 只处理HTTP请求
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return
        
        # 检查是否需要缓存请求体
        method = scope.get("method", "")
        if method not in ["POST", "PUT", "PATCH"]:
            # GET、DELETE等请求通常没有请求体，直接传递
            await self.app(scope, receive, send)
            return
        
        # 读取并缓存请求体
        try:
            body, new_receive = await self._cache_request_body(receive)
            
            # 将缓存的请求体存储在scope中，供后续使用
            scope.setdefault("state", {})
            scope["state"]["cached_body"] = body
            
            logger.debug(f"请求体已缓存: {len(body)} 字节")
            
            # 使用新的receive函数调用应用
            await self.app(scope, new_receive, send)
            
        except Exception as e:
            logger.error(f"请求体缓存失败: {e}")
            # 出错时回退到原始行为
            await self.app(scope, receive, send)
    
    async def _cache_request_body(self, receive: Receive) -> Tuple[bytes, Receive]:
        """
        缓存请求体并创建新的receive函数
        
        Args:
            receive: 原始的receive函数
            
        Returns:
            Tuple[bytes, Receive]: (缓存的请求体, 新的receive函数)
        """
        body_parts = []
        body_size = 0
        
        async def receive_body() -> Message:
            """读取请求体消息"""
            nonlocal body_size
            
            message = await receive()
            
            if message["type"] == "http.request":
                body_chunk = message.get("body", b"")
                body_size += len(body_chunk)
                
                # 检查请求体大小限制
                if body_size > self.max_body_size:
                    raise ValueError(f"请求体超过大小限制: {self.max_body_size} 字节")
                
                body_parts.append(body_chunk)
            
            return message
        
        # 读取完整的请求体
        while True:
            message = await receive_body()
            
            if message["type"] == "http.request":
                # 检查是否还有更多数据
                if not message.get("more_body", False):
                    break
            elif message["type"] == "http.disconnect":
                # 客户端断开连接
                break
            else:
                # 其他类型的消息，继续处理
                continue
        
        # 合并所有请求体部分
        complete_body = b"".join(body_parts)
        
        # 创建新的receive函数，可以重复返回缓存的请求体
        body_sent = False
        
        async def new_receive() -> Message:
            """新的receive函数，返回缓存的请求体"""
            nonlocal body_sent
            
            if not body_sent:
                body_sent = True
                return {
                    "type": "http.request",
                    "body": complete_body,
                    "more_body": False,
                }
            else:
                # 请求体已发送，后续调用返回空
                return {
                    "type": "http.request", 
                    "body": b"",
                    "more_body": False,
                }
        
        return complete_body, new_receive


def get_cached_body(scope: Scope) -> bytes:
    """
    从scope中获取缓存的请求体
    
    Args:
        scope: ASGI scope
        
    Returns:
        bytes: 缓存的请求体，如果没有则返回空字节
    """
    state = scope.get("state", {})
    return state.get("cached_body", b"")


def setup_request_body_cache(request) -> bytes:
    """
    从FastAPI Request对象中获取缓存的请求体
    
    Args:
        request: FastAPI Request对象
        
    Returns:
        bytes: 缓存的请求体
    """
    if hasattr(request, "scope"):
        return get_cached_body(request.scope)
    return b""